#include "Util.h"
#include "Error.h"

#include <urlmon.h>
#include <iphlpapi.h>
#include <WinInet.h>

#pragma comment(lib, "urlmon.lib")
#pragma comment(lib, "iphlpapi")

bool Util::DownloadFile( __in std::string URL, __in std::string Name )
{
	if( URLDownloadToFile( NULL, URL.c_str( ), Name.c_str( ), 0, NULL ) != S_OK )
		return false;

	return true;
}

BOOL Util::EnableDebugPrivilege( bool Enable )
{
	BOOL Result = FALSE;
	HANDLE TokenHandle = NULL;
	TOKEN_PRIVILEGES TokenPrivileges;

	do
	{
		if( !OpenProcessToken( GetCurrentProcess( ), TOKEN_ADJUST_PRIVILEGES, &TokenHandle ) )
		{
			Result = ERROR_FAILED_TO_OPEN_PROCESS_TOKEN;
			break;
		}
		
		TokenPrivileges.PrivilegeCount = 1;
		
		if( !LookupPrivilegeValue( NULL, SE_DEBUG_NAME, &TokenPrivileges.Privileges[ 0 ].Luid ) )
		{
			Result = ERROR_FAILED_TO_LOOKUP_PRIVILEGE_VALUE;
            break;
        }
		
		TokenPrivileges.Privileges[ 0 ].Attributes = Enable ? SE_PRIVILEGE_ENABLED : 0;
        
		if( !AdjustTokenPrivileges( TokenHandle, FALSE, &TokenPrivileges, sizeof( TokenPrivileges ), NULL, NULL ) )
		{
			Result = ERROR_FAILED_TO_ADJUST_TOKEN_PRIVILEGE;
			break;
		}

		Result = TRUE;
	} while ( 0 );

	CloseHandle( TokenHandle );
	return Result;
}

std::string Util::GetCurDirectory( void )
{
	char Buffer[ MAX_PATH ] = { NULL };
	GetModuleFileName( NULL, Buffer, sizeof( Buffer ) );

	return std::string( Buffer ).substr( 0, std::string( Buffer ).find_last_of( "\\/" ) + 1 );
}